from run import app
import arrow
import json
from app.db.models.jobs import Job, CustomSecurityEventSettings, CustomSslSettings, CustomVulSettings
from app.db.models.tasks import Task
from app.db.models.assets import CustomTaskSettings
from app.libs import scanner
from app.libs.utils import schedule_setting

task_ids = []
send = False
success_task_ids = set()
error_task_ids = set()

# 测试时修改时间, 注意是utc时间
now = '2022-08-10T05:18:56.000000+00:00'
now = arrow.get(now)
now = arrow.Arrow.utcnow()

info = f"""
###############################################
开始执行tasks未完成的任务迁移...
是否发送任务: {send}
当前时间: {now.for_json()}
###############################################
"""
print(info)

if task_ids:
    tasks = Task.objects.filter(pk__in=task_ids)
else:
    tasks = Task.objects.filter(status__in=['waiting', 'active', 'stalled'])

info = f"""
需要处理的tasks数量: {tasks.count()}"""
print(info)

SCHEDULER_MAP = {
    'asset': scanner.AssetScheduler,
    'securityEvent': scanner.SecurityEventScheduler,
    'ssl': scanner.SslScheduler,
    'vul': scanner.VulScheduler,
}

SETTING_MAP = {
    'asset': CustomTaskSettings,
    'securityEvent': CustomSecurityEventSettings,
    'ssl': CustomSslSettings,
    'vul': CustomVulSettings
}

for task in tasks:
    print(
        f"domain: {task.domain}, task_type: {task.taskType}, triggerType:{task.triggerType}, startTime:{task.startTime} status: {task.status}")
    if task.taskType in ['http', 'ping']:
        print("可用性, 删除")
        if send:
            task.delete()
        continue

    if arrow.Arrow.fromdatetime(task.startTime) > now:
        print("未来的任务,删除")
        if send:
            task.delete()
        continue

    job = Job.objects.with_id(task.jobId)
    setting = SETTING_MAP[task.taskType].from_dict(task.taskSettings.to_dict())
    addition = {'pk': str(task.pk)}
    if send:
        try:
            task_id = schedule_setting(
                SCHEDULER_MAP[task.taskType],
                job,
                setting,
                trigger_type=task.triggerType,
                immediate_exec=True,
                index=task.taskIdx,
                addition=addition,
                is_periodic=False,
                is_disabled=False
            )
            success_task_ids.add(task_id)
        except Exception as e:
            error_task_ids.add(str(task.pk))

task_id_groups = []
for task_id in success_task_ids:
    task_id_groups.append({
        'task_id': task_id
    })
print(f"\n删除调度任务需要的参数:\n{json.dumps(task_id_groups, indent=4)}")

print(f"\n失败的task_id:\n{error_task_ids}")
